// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <benchmark/benchmark.h>

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <memory>
#include <random>
#include <string>
#include <vector>

#include <nebula/core/buffer.h>
#include <nebula/io/compressed.h>
#include <nebula/io/memory.h>

#include <nebula/testing/ktest_util.h>
#include <nebula/compression/compression.h>
#include <nebula/version.h>
#include <turbo/log/logging.h>
#include <turbo/base/macros.h>

namespace nebula::io {

using ::nebula::Compression;

std::vector<uint8_t> MakeCompressibleData(int data_size) {
  // XXX This isn't a real-world corpus so doesn't really represent the
  // comparative qualities of the algorithms

  // First make highly compressible data
  std::string base_data =
      "Apache Nebula is a cross-language development platform for in-memory data";
  int nrepeats = static_cast<int>(1 + data_size / base_data.size());

  std::vector<uint8_t> data(base_data.size() * nrepeats);
  for (int i = 0; i < nrepeats; ++i) {
    std::memcpy(data.data() + i * base_data.size(), base_data.data(), base_data.size());
  }
  data.resize(data_size);

  // Then randomly mutate some bytes so as to make things harder
  std::mt19937 engine(42);
  std::exponential_distribution<> offsets(0.05);
  std::uniform_int_distribution<> values(0, 255);

  int64_t pos = 0;
  while (pos < data_size) {
    data[pos] = static_cast<uint8_t>(values(engine));
    pos += static_cast<int64_t>(offsets(engine));
  }

  return data;
}

// Using a non-zero copy buffer reader to benchmark the non-zero copy path.
class NonZeroCopyBufferReader final : public InputStream {
 public:
  NonZeroCopyBufferReader(std::shared_ptr<Buffer> buffer) : reader_(std::move(buffer)) {}

  bool supports_zero_copy() const override { return false; }

  turbo::Result<int64_t> read(int64_t nbytes, void* out) override {
    return reader_.read(nbytes, out);
  }

  turbo::Result<std::shared_ptr<Buffer>> read(int64_t nbytes) override {
    // Testing the non-zero copy path like reading from local file or Object store,
    // so we need to allocate a buffer and copy the data.
    TURBO_MOVE_OR_RAISE(auto buf, ::nebula::allocate_resizable_buffer(nbytes));
    TURBO_MOVE_OR_RAISE(int64_t size, read(nbytes, buf->mutable_data()));
    TURBO_RETURN_NOT_OK(buf->resize(size));
    return buf;
  }
  turbo::Status close() override { return reader_.close(); }
  turbo::Result<int64_t> tell() const override { return reader_.tell(); }
  bool closed() const override { return reader_.closed(); }

 private:
  ::nebula::io::BufferReader reader_;
};

enum class BufferReadMode { ProvidedByCaller, ReturnedByCallee };

template <typename BufReader, BufferReadMode Mode>
static void CompressedInputStreamBenchmark(::benchmark::State& state,
                                           CompressionType compression) {
  const int64_t input_size = state.range(0);
  const int64_t batch_size = state.range(1);

  const std::vector<uint8_t> data = MakeCompressibleData(static_cast<int>(input_size));
  auto codec = ::nebula::Codec::create(compression).value_or_die();
  int64_t max_compress_len =
      codec->MaxCompressedLen(static_cast<int64_t>(data.size()), data.data());
  std::shared_ptr<::nebula::ResizableBuffer> buf =
      ::nebula::allocate_resizable_buffer(max_compress_len).value_or_die();
  const int64_t compressed_length =
      codec
          ->Compress(static_cast<int64_t>(data.size()), data.data(), max_compress_len,
                     buf->mutable_data())
          .value_or_die();
  ABORT_NOT_OK(buf->resize(compressed_length));
  for (auto _ : state) {
    state.PauseTiming();
    auto reader = std::make_shared<BufReader>(buf);
    [[maybe_unused]] std::unique_ptr<Buffer> read_buffer;
    if constexpr (Mode == BufferReadMode::ProvidedByCaller) {
      read_buffer = ::nebula::allocate_buffer(batch_size).value_or_die();
    }
    state.ResumeTiming();
    // Put `CompressedInputStream::create` in timing.
    auto input_stream =
        ::nebula::io::CompressedInputStream::create(codec.get(), reader).value_or_die();
    auto remaining_size = input_size;
    while (remaining_size > 0) {
      if constexpr (Mode == BufferReadMode::ProvidedByCaller) {
        auto value = input_stream->read(batch_size, read_buffer->mutable_data());
        ABORT_NOT_OK(value);
        remaining_size -= value.value_or_die();
      } else {
        auto value = input_stream->read(batch_size);
        ABORT_NOT_OK(value);
        remaining_size -= value.value_or_die()->size();
      }
    }
  }
  state.SetBytesProcessed(input_size * state.iterations());
}

template <CompressionType kCompression>
static void CompressedInputStreamZeroCopyBufferProvidedByCaller(
    ::benchmark::State& state) {
  CompressedInputStreamBenchmark<::nebula::io::BufferReader,
                                 BufferReadMode::ProvidedByCaller>(state, kCompression);
}

template <CompressionType kCompression>
static void CompressedInputStreamNonZeroCopyBufferProvidedByCaller(
    ::benchmark::State& state) {
  CompressedInputStreamBenchmark<NonZeroCopyBufferReader,
                                 BufferReadMode::ProvidedByCaller>(state, kCompression);
}

template <CompressionType kCompression>
static void CompressedInputStreamZeroCopyBufferReturnedByCallee(
    ::benchmark::State& state) {
  CompressedInputStreamBenchmark<::nebula::io::BufferReader,
                                 BufferReadMode::ReturnedByCallee>(state, kCompression);
}

template <CompressionType kCompression>
static void CompressedInputStreamNonZeroCopyBufferReturnedByCallee(
    ::benchmark::State& state) {
  CompressedInputStreamBenchmark<NonZeroCopyBufferReader,
                                 BufferReadMode::ReturnedByCallee>(state, kCompression);
}

static void CompressedInputArguments(::benchmark::internal::Benchmark* b) {
  b->ArgNames({"num_bytes", "batch_size"})
      ->Args({8 * 1024, 8 * 1024})
      ->Args({64 * 1024, 8 * 1024})
      ->Args({64 * 1024, 64 * 1024})
      ->Args({1024 * 1024, 8 * 1024})
      ->Args({1024 * 1024, 64 * 1024})
      ->Args({1024 * 1024, 1024 * 1024});
}

// Benchmark LZ4 because it's lightweight, which makes benchmarking focused on the
// overhead of the compression input stream.
BENCHMARK_TEMPLATE(CompressedInputStreamZeroCopyBufferProvidedByCaller,
                   CompressionType::LZ4_FRAME)
    ->Apply(CompressedInputArguments);
BENCHMARK_TEMPLATE(CompressedInputStreamNonZeroCopyBufferProvidedByCaller,
                   CompressionType::LZ4_FRAME)
    ->Apply(CompressedInputArguments);
BENCHMARK_TEMPLATE(CompressedInputStreamZeroCopyBufferReturnedByCallee,
                   CompressionType::LZ4_FRAME)
    ->Apply(CompressedInputArguments);
BENCHMARK_TEMPLATE(CompressedInputStreamNonZeroCopyBufferReturnedByCallee,
                   CompressionType::LZ4_FRAME)
    ->Apply(CompressedInputArguments);

}  // namespace nebula::io
